#include <ros/ros.h>
#include "mcu_correspond/mcu_data_subscriber.h"
#include <opencv2/opencv.hpp>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/Image.h>
#include <cmath>
#include <locale>

class ColorDetector {
private:
    ros::NodeHandle nh_;
    ros::Subscriber image_sub_;
    ros::Publisher image_pub_;
    McuDataSubscriber mcu_data_;
    ros::Rate publish_rate_;
    ros::Time last_publish_time_;

public:
    ColorDetector(ros::NodeHandle& nh) : nh_(nh), mcu_data_(nh), publish_rate_(1.0) {
        // Subscribe to the camera image topic
        image_sub_ = nh_.subscribe("/usb_cam/image_raw", 1, &ColorDetector::imageCallback, this);
        // Publish the processed image
        image_pub_ = nh_.advertise<sensor_msgs::Image>("/processed_image", 1);
        
        last_publish_time_ = ros::Time::now();
        ROS_INFO("HSV Color Detector initialized. Detecting Red and Green areas.");
    }
    
    void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
        try {
            // Convert ROS image message to OpenCV format and resize
            cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
            cv::Mat bgr_image;
            cv::resize(cv_ptr->image, bgr_image, cv::Size(640, 360));

            // Convert BGR to HSV
            cv::Mat hsv_image;
            cv::cvtColor(bgr_image, hsv_image, cv::COLOR_BGR2HSV);

            // Define HSV ranges for Red color
            // Red wraps around 180, so we need two ranges
            cv::Scalar lower_red1(0, 70, 50);
            cv::Scalar upper_red1(10, 255, 255);
            cv::Scalar lower_red2(170, 70, 50);
            cv::Scalar upper_red2(180, 255, 255);

            // Define HSV range for Green color
            cv::Scalar lower_green(35, 70, 50);
            cv::Scalar upper_green(85, 255, 255);

            // Create masks for red and green
            cv::Mat red_mask1, red_mask2, red_mask;
            cv::inRange(hsv_image, lower_red1, upper_red1, red_mask1);
            cv::inRange(hsv_image, lower_red2, upper_red2, red_mask2);
            cv::add(red_mask1, red_mask2, red_mask); // Combine the two red masks

            cv::Mat green_mask;
            cv::inRange(hsv_image, lower_green, upper_green, green_mask);

            // Calculate the area of red and green regions for the whole image
            int total_red_area = cv::countNonZero(red_mask);
            int total_green_area = cv::countNonZero(green_mask);

            // Log the total area sizes
            ROS_INFO("Total Area - Red: %d, Green: %d pixels", total_red_area, total_green_area);

            // --- Center Grid Analysis (九宫格中间一格) ---
            // Define the center grid cell ROI
            int start_x = bgr_image.cols / 3;
            int end_x = 2 * bgr_image.cols / 3;
            int start_y = bgr_image.rows / 3;
            int end_y = 2 * bgr_image.rows / 3;
            cv::Rect center_roi(start_x, start_y, end_x - start_x, end_y - start_y);

            // Get the color masks for the center ROI
            cv::Mat center_red_mask = red_mask(center_roi);
            cv::Mat center_green_mask = green_mask(center_roi);

            // Calculate the area of red and green regions in the center
            int center_red_area = cv::countNonZero(center_red_mask);
            int center_green_area = cv::countNonZero(center_green_mask);

            // Determine and log the dominant color in the center grid
            std::string dominant_color = "None";
            if (center_red_area > center_green_area) {
                dominant_color = "Red";
            } else if (center_green_area > center_red_area) {
                dominant_color = "Green";
            } else if (center_red_area > 0) { // Equal and not zero
                dominant_color = "Equal";
            }
            ROS_INFO("Center Grid - Dominant: %s (Red: %d, Green: %d)", dominant_color.c_str(), center_red_area, center_green_area);

            if (dominant_color == "Red" || dominant_color == "Green") {
                // Publish the dominant color to the MCU at 1 Hz
                ros::Time current_time = ros::Time::now();
                if (current_time - last_publish_time_ >= ros::Duration(1.0)) {
                    mcu_data_.publishDoubleData("rog", dominant_color == "Red" ? 1.0 : 2.0);
                    last_publish_time_ = current_time;
                }
            }
            else {
                // Publish the dominant color to the MCU at 1 Hz
                ros::Time current_time = ros::Time::now();
                if (current_time - last_publish_time_ >= ros::Duration(1.0)) {
                    mcu_data_.publishDoubleData("rog", 0.0);
                    last_publish_time_ = current_time;
                }
            }
            

            // Create a result image for visualization
            cv::Mat result_image = cv::Mat::zeros(bgr_image.size(), CV_8UC3);
            result_image.setTo(cv::Scalar(0, 0, 255), red_mask);   // Mark red areas in red
            result_image.setTo(cv::Scalar(0, 255, 0), green_mask); // Mark green areas in green
            
            // Publish the processed image
            sensor_msgs::ImagePtr output_msg = cv_bridge::CvImage(std_msgs::Header(), "bgr8", result_image).toImageMsg();
            image_pub_.publish(output_msg);
            
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
        }
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "s2020basic");
    ros::NodeHandle nh;
    
    // Create color detector
    ColorDetector detector(nh);
    
    ros::spin();
    return 0;
}
